"""
文档加载器

支持加载多种格式的文档：PDF、Word、TXT、网页等
"""

import os
import logging
from typing import List, Dict, Any, Optional
from pathlib import Path

from langchain_community.document_loaders import (
    PyPDFLoader,
    Docx2txtLoader,
    TextLoader,
    WebBaseLoader,
    CSVLoader,
    UnstructuredMarkdownLoader
)
from langchain_core.documents import Document

logger = logging.getLogger(__name__)

class DocumentLoader:
    """
    文档加载器类
    
    支持多种文档格式的加载和处理
    """
    
    def __init__(self):
        """初始化文档加载器"""
        self.supported_formats = {
            '.pdf': self._load_pdf,
            '.docx': self._load_docx,
            '.doc': self._load_docx,
            '.txt': self._load_txt,
            '.md': self._load_markdown,
            '.csv': self._load_csv
        }
    
    def load_file(self, file_path: str) -> List[Document]:
        """
        加载单个文件
        
        Args:
            file_path (str): 文件路径
            
        Returns:
            List[Document]: 加载的文档列表
        """
        try:
            file_path = Path(file_path)
            
            if not file_path.exists():
                raise FileNotFoundError(f"文件不存在: {file_path}")
            
            suffix = file_path.suffix.lower()
            
            if suffix not in self.supported_formats:
                raise ValueError(f"不支持的文件格式: {suffix}")
            
            logger.info(f"正在加载文件: {file_path}")
            documents = self.supported_formats[suffix](str(file_path))
            
            # 添加元数据
            for doc in documents:
                doc.metadata.update({
                    'source': str(file_path),
                    'file_name': file_path.name,
                    'file_size': file_path.stat().st_size,
                    'file_type': suffix
                })
            
            logger.info(f"成功加载 {len(documents)} 个文档块")
            return documents
            
        except Exception as e:
            logger.error(f"加载文件失败 {file_path}: {str(e)}")
            raise
    
    def load_directory(self, directory_path: str, recursive: bool = True) -> List[Document]:
        """
        加载目录中的所有支持的文档
        
        Args:
            directory_path (str): 目录路径
            recursive (bool): 是否递归加载子目录
            
        Returns:
            List[Document]: 加载的文档列表
        """
        try:
            directory = Path(directory_path)
            
            if not directory.exists() or not directory.is_dir():
                raise ValueError(f"目录不存在或不是有效目录: {directory_path}")
            
            all_documents = []
            
            # 获取文件列表
            if recursive:
                files = directory.rglob('*')
            else:
                files = directory.iterdir()
            
            # 过滤支持的文件格式
            supported_files = [
                f for f in files 
                if f.is_file() and f.suffix.lower() in self.supported_formats
            ]
            
            logger.info(f"找到 {len(supported_files)} 个支持的文件")
            
            # 逐个加载文件
            for file_path in supported_files:
                try:
                    documents = self.load_file(str(file_path))
                    all_documents.extend(documents)
                except Exception as e:
                    logger.warning(f"跳过文件 {file_path}: {str(e)}")
                    continue
            
            logger.info(f"总共加载了 {len(all_documents)} 个文档块")
            return all_documents
            
        except Exception as e:
            logger.error(f"加载目录失败 {directory_path}: {str(e)}")
            raise
    
    def load_urls(self, urls: List[str]) -> List[Document]:
        """
        加载网页内容
        
        Args:
            urls (List[str]): 网页URL列表
            
        Returns:
            List[Document]: 加载的文档列表
        """
        try:
            all_documents = []
            
            for url in urls:
                try:
                    logger.info(f"正在加载网页: {url}")
                    loader = WebBaseLoader(url)
                    documents = loader.load()
                    
                    # 添加元数据
                    for doc in documents:
                        doc.metadata.update({
                            'source': url,
                            'source_type': 'web',
                            'url': url
                        })
                    
                    all_documents.extend(documents)
                    logger.info(f"成功加载网页，获得 {len(documents)} 个文档块")
                    
                except Exception as e:
                    logger.warning(f"加载网页失败 {url}: {str(e)}")
                    continue
            
            return all_documents
            
        except Exception as e:
            logger.error(f"加载网页失败: {str(e)}")
            raise
    
    def _load_pdf(self, file_path: str) -> List[Document]:
        """加载PDF文件"""
        loader = PyPDFLoader(file_path)
        return loader.load()
    
    def _load_docx(self, file_path: str) -> List[Document]:
        """加载Word文档"""
        loader = Docx2txtLoader(file_path)
        return loader.load()
    
    def _load_txt(self, file_path: str) -> List[Document]:
        """加载文本文件"""
        loader = TextLoader(file_path, encoding='utf-8')
        return loader.load()
    
    def _load_markdown(self, file_path: str) -> List[Document]:
        """加载Markdown文件"""
        loader = UnstructuredMarkdownLoader(file_path)
        return loader.load()
    
    def _load_csv(self, file_path: str) -> List[Document]:
        """加载CSV文件"""
        loader = CSVLoader(file_path)
        return loader.load()
    
    def get_document_info(self, documents: List[Document]) -> Dict[str, Any]:
        """
        获取文档信息统计
        
        Args:
            documents (List[Document]): 文档列表
            
        Returns:
            Dict[str, Any]: 文档信息统计
        """
        if not documents:
            return {
                'total_documents': 0,
                'total_characters': 0,
                'file_types': {},
                'sources': []
            }
        
        total_chars = sum(len(doc.page_content) for doc in documents)
        file_types = {}
        sources = set()
        
        for doc in documents:
            # 统计文件类型
            file_type = doc.metadata.get('file_type', 'unknown')
            file_types[file_type] = file_types.get(file_type, 0) + 1
            
            # 收集来源
            source = doc.metadata.get('source', 'unknown')
            sources.add(source)
        
        return {
            'total_documents': len(documents),
            'total_characters': total_chars,
            'average_length': total_chars // len(documents) if documents else 0,
            'file_types': file_types,
            'sources': list(sources),
            'unique_sources': len(sources)
        }
    
    def get_supported_formats(self) -> List[str]:
        """获取支持的文件格式列表"""
        return list(self.supported_formats.keys())

# 创建全局实例
document_loader = DocumentLoader() 